#!/usr/bin/env python

""" display_image.py - Version 0.1 2011-06-08

    Display an OpenCV image that is part of the pi_rein pipeline.
    
    Created for the Pi Robot Project: http://www.pirobot.org
    Copyright (c) 2011 Patrick Goebel.  All rights reserved.

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.
    
    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details at:
    
    http://www.gnu.org/licenses/gpl.html
      
"""

import roslib
roslib.load_manifest('pi_rein')
import rospy
import cv
import sys
from pi_rein.msg import DetectionArray
from pi_rein.srv import *
from sensor_msgs.msg import Image, CameraInfo
from cv_bridge import CvBridge, CvBridgeError
import time
import message_filters


class DisplayImage:
    def __init__(self, node_name):
        rospy.init_node(node_name)
        
        rospy.on_shutdown(self.cleanup)
        
        self.node_name = node_name
        self.show_text = rospy.get_param('~show_text', True)
        self.show_labels = rospy.get_param('~show_labels', True)
        self.show_features = rospy.get_param('~show_features', True)
        self.flip_image = rospy.get_param('~flip_image', False)
        self.marker_topics = rospy.get_param('~marker_topics', {'detections': DetectionArray, 'random_markers': DetectionArray})
        
        """ Initialize a number of global variables """
        self.image = None
        self.marker_image = None
        self.display_image = None
        self.keystroke = None
        self.detect_box = None
        self.track_box = None
        self.display_box = None
        self.keep_marker_history = False
        self.night_mode = False
        self.cps = 0 # Cylces per second = number of processing loops per second.
        self.cps_values = list()
        self.cps_n_values = 20

        """ Create the display window """
        self.cv_window_name = self.node_name
        cv.NamedWindow(self.cv_window_name, cv.CV_NORMAL)
        cv.ResizeWindow(self.cv_window_name, 640, 480)
        
        """ Create the cv_bridge object """
        self.cv_bridge = CvBridge()
    
        subscribers = []
        
        image_sub = message_filters.Subscriber("image", Image)
        
        subscribers.append(image_sub)
        
        for topic, type in self.marker_topics.iteritems():
            topic_sub = message_filters.Subscriber(topic, type)
            subscribers.append(topic_sub)
            
        #""" Subscribe to the raw camera image topic and set the image processing callback """
        #self.image_sub = message_filters.Subscriber("image", Image)
        
        #""" Subscribe to a topic that provides markers """
        #self.detections_sub = message_filters.Subscriber("detections", DetectionArray)
        
        #tss = message_filters.TimeSynchronizer([self.image_sub, self.detections_sub], 100)
        tss = message_filters.TimeSynchronizer(subscribers, 100)
        tss.registerCallback(self.sync_callback)

        rospy.loginfo("Starting " + self.node_name)
        
        rospy.wait_for_message("image", Image, timeout=15)
        
    def sync_callback(self, image, detections, markers):
        """ Time this loop to get cycles per second """
        start = time.time()
        
        self.image_callback(image)
        self.detections_callback(detections)
        self.detections_callback(markers)
        
        end = time.time()
        duration = end - start
        fps = int(1.0 / duration)
        self.cps_values.append(fps)
        if len(self.cps_values) > self.cps_n_values:
            self.cps_values.pop(0)
        self.cps = int(sum(self.cps_values) / len(self.cps_values))
    
        
    def detections_callback(self, detections):
        #if not self.keep_marker_history:
            #cv.Zero(self.marker_image)
        for detection in detections.detections:
            pt1 = (detection.mask.roi.x, detection.mask.roi.y)
            pt2 = (detection.mask.roi.x + detection.mask.roi.width, detection.mask.roi.y + detection.mask.roi.height)
            cv.Rectangle(self.marker_image, pt1, pt2, cv.RGB(255, 0, 0), 2, 8, 0)
 
    def image_callback(self, ros_image):
        """ Convert the raw image to OpenCV format """
        (self.image, self.image_size, self.image_width, self.image_height) = self.ros2cv_image(ros_image, "bgr8")
        
        """ Some webcams invert the image """
        if self.flip_image:
            cv.Flip(self.image)
                    
        """ Create a few images we will use for display """
        if not self.display_image:
            self.display_image = cv.CreateImage(self.image_size, 8, 3)
            self.marker_image = cv.CreateImage(self.image_size, 8, 3)
        
        if self.night_mode:
            """ Night mode: only display the markers """
            cv.SetZero(self.image)
        
        """ Merge the raw image and the marker image to form the display image """
        cv.Or(self.image, self.marker_image, self.display_image)
        #cv.Copy(self.image, self.display_image)
        
        if self.show_text:
            text_font = cv.InitFont(cv.CV_FONT_VECTOR0, 1, 1, 0, 2, 8)
            """ Print cycles per second (CPS) and resolution (RES) at top of the image """
            cv.PutText(self.display_image, "CPS: " + str(self.cps), (10, int(self.image_size[1] * 0.1)), text_font, cv.RGB(255, 255, 0))
            #cv.PutText(self.display_image, "RES: " + str(self.image_size[0]) + "X" + str(self.image_size[1]), (int(self.image_size[0] * 0.6), int(self.image_size[1] * 0.1)), text_font, cv.RGB(255, 255, 0))

        """ Show the final display image """
        cv.ShowImage(self.cv_window_name, self.display_image)
        
        """ Look for keyboard input from the user """
        keystroke = cv.WaitKey(6)
        
        """ Process any keyboard commands """
        if 32 <= self.keystroke and self.keystroke < 128:
            cc = chr(self.keystroke).lower()
            if cc == 'c':
                self.features = []
                self.track_box = None
                self.detect_box = None
            elif cc == 'n':
                self.night_mode = not self.night_mode
            elif cc == 'f':
                self.show_features = not self.show_features
            elif cc == 't':
                self.show_text = not self.show_text
            elif cc == 'q':
                """ user has press the q key, so exit """
                rospy.signal_shutdown("User hit q key to quit.")
      
    def ros2cv_image(self, ros_image, type):
        try:
            cv_image = self.cv_bridge.imgmsg_to_cv(ros_image, type)
            size = cv.GetSize(cv_image)
            width = size[0]
            height = size[1]
            return (cv_image, size, width, height)
        except CvBridgeError, e:
            print e
            
            
#    def handle_mouse(self):
#        pass
            
        # If the user is selecting a region with the mouse, display the corresponding rectangle for feedback.
#        if self.drag_start and self.is_rect_nonzero(self.selection):
#            x,y,w,h = self.selection
#            cv.Rectangle(self.marker_image, (x, y), (x + w, y + h), (0, 255, 255), 2)
#            self.selected_point = None
#
#        # Else if the user has clicked on a point on the image, display it as a small circle.            
#        elif not self.selected_point is None:
#            x = self.selected_point[0]
#            y = self.selected_point[1]
#            cv.Circle(self.marker_image, (x, y), 3, (0, 255, 255), 2)
        
#    def is_rect_nonzero(self, r):
#        # First assume a simple CvRect type
#        try:
#            (_,_,w,h) = r
#            return (w > 0) and (h > 0)
#        except:
#            # Otherwise, assume a CvBox2D type
#            ((_,_),(w,h),a) = r
#            return (w > 0) and (h > 0) 
        
    def cleanup(self):
        print "Shutting down vision node."
        cv.DestroyAllWindows()       

def main(args):
    # Display a help message if appropriate.
    help_message = ""
          
    print help_message
    
    try:
        DisplayImage("display_image")
        rospy.spin()
    except KeyboardInterrupt:
        print "Shutting down vision node."
        cv.DestroyAllWindows()

if __name__ == '__main__':
    main(sys.argv)
